import { BaseEmbeddings } from '@cherrystudio/embedjs-interfaces'

export interface GlobalSemanticSplitterConfig {
  minChunkSize?: number
  similarityThreshold?: number
}

export class GlobalSemanticSplitter {
  private readonly embeddingModel: BaseEmbeddings
  private readonly config: Required<GlobalSemanticSplitterConfig>

  constructor(embeddingModel: BaseEmbeddings, config: GlobalSemanticSplitterConfig = {}) {
    this.embeddingModel = embeddingModel
    this.config = {
      minChunkSize: 100,
      similarityThreshold: 0.7,
      ...config
    }
  }

  private splitIntoSentences(text: string): string[] {
    // ...同SemanticSplitter实现...
    text = text.replace(/ /g, '') // 移除多余空格
    const rawSentences = text.split(/(?<=[。！？!?])\s*/).map((s) => s.trim())
    const mergedSentences: string[] = []
    for (let i = 0; i < rawSentences.length; i++) {
      const current = rawSentences[i]
      if (current.length < 10 && i < rawSentences.length - 1) {
        rawSentences[i + 1] = current + rawSentences[i + 1]
      } else {
        mergedSentences.push(current)
      }
    }
    return mergedSentences
  }

  private async getEmbeddingsBatch(sentences: string[]): Promise<number[][]> {
    if (typeof (this.embeddingModel as any).embedDocuments === 'function') {
      return await (this.embeddingModel as any).embedDocuments(sentences)
    } else {
      return await Promise.all(sentences.map((s) => this.embeddingModel.embedQuery(s)))
    }
  }

  /**
   * 简单KMeans聚类实现
   * @param vectors 向量数组
   * @param k 聚类数
   * @param maxIter 最大迭代次数
   * @returns 每个向量的聚类标签
   */
  private kmeans(vectors: number[][], k: number, maxIter = 20): number[] {
    // 随机初始化k个中心
    const centers = vectors.slice(0, k).map((v) => v.slice())
    const labels = new Array(vectors.length).fill(0)
    for (let iter = 0; iter < maxIter; iter++) {
      // 分配标签
      for (let i = 0; i < vectors.length; i++) {
        let best = 0
        let bestDist = this.euclideanDist(vectors[i], centers[0])
        for (let j = 1; j < k; j++) {
          const dist = this.euclideanDist(vectors[i], centers[j])
          if (dist < bestDist) {
            best = j
            bestDist = dist
          }
        }
        labels[i] = best
      }
      // 更新中心
      const sums = Array.from({ length: k }, () => new Array(vectors[0].length).fill(0))
      const counts = new Array(k).fill(0)
      for (let i = 0; i < vectors.length; i++) {
        const label = labels[i]
        for (let d = 0; d < vectors[i].length; d++) {
          sums[label][d] += vectors[i][d]
        }
        counts[label]++
      }
      for (let j = 0; j < k; j++) {
        if (counts[j] === 0) continue
        for (let d = 0; d < sums[j].length; d++) {
          centers[j][d] = sums[j][d] / counts[j]
        }
      }
    }
    return labels
  }

  private euclideanDist(a: number[], b: number[]): number {
    let sum = 0
    for (let i = 0; i < a.length; i++) {
      const diff = a[i] - b[i]
      sum += diff * diff
    }
    return Math.sqrt(sum)
  }

  /**
   * 基于KMeans聚类的全局语义分片
   */
  async splitText(text: string): Promise<string[]> {
    const sentences = this.splitIntoSentences(text)
    if (sentences.length === 0) return []

    const embeddings = await this.getEmbeddingsBatch(sentences)

    // 1. 估算聚类数k（可根据文本长度自适应，也可固定）
    let k = Math.max(1, Math.floor(sentences.join('').length / this.config.minChunkSize))
    k = Math.min(k, sentences.length)

    // 2. KMeans聚类
    const labels = this.kmeans(embeddings, k)

    // 3. 按聚类标签分组，保持原顺序
    const clusterMap: Map<number, { idx: number; text: string }[]> = new Map()
    for (let i = 0; i < labels.length; i++) {
      if (!clusterMap.has(labels[i])) clusterMap.set(labels[i], [])
      clusterMap.get(labels[i])!.push({ idx: i, text: sentences[i] })
    }
    // 按原始顺序合并同一聚类的句子
    const clusters = Array.from(clusterMap.values()).map((arr) => arr.sort((a, b) => a.idx - b.idx).map((x) => x.text))

    // 4. 合并chunk以满足最小尺寸
    const chunks: string[] = []
    let buffer = ''
    let bufferLen = 0
    for (let i = 0; i < clusters.length; i++) {
      const chunkText = clusters[i].join(' ')
      if (bufferLen + chunkText.length < this.config.minChunkSize && i < clusters.length - 1) {
        buffer += (buffer ? ' ' : '') + chunkText
        bufferLen += chunkText.length
      } else {
        const finalChunk = buffer ? buffer + ' ' + chunkText : chunkText
        chunks.push(finalChunk)
        buffer = ''
        bufferLen = 0
      }
    }
    if (buffer) {
      if (chunks.length > 0) {
        chunks[chunks.length - 1] += ' ' + buffer
      } else {
        chunks.push(buffer)
      }
    }
    return chunks
  }
}
